#include "gc.h"

void debug_graph_connection (Graph3D const& graph, int startSliceNo, 
			     int layerNum);
int greedy (Graph3D& graph, std::string const& initialFilePath, 
	    std::vector<int>& targetStartComponentNoSet, 
	    int startSliceNo, int startLayerNo, int layerNum, 
	    SegmentationGraph& fwdSegGraph, SegmentationGraph& bwdSegGraph, 
	    double& fwdMaxFlow, double& bwdMaxFlow, 
	    VerboseLevel verbose = VERBOSE_LOW);
int greedy (Graph3D& graph, std::string const& initialFilePath, 
	    int startSliceNo, int startLayerNo, int layerNum, 
	    VerboseLevel verbose = VERBOSE_LOW);

void initialization (int argc, char* argv[], 
		     VerboseLevel& verbose, 
		     double& sigma_area, double& sigma_correlation, 
		     double& alpha, 
		     int& startSliceNo, int& endSliceNo, int& pileLayerNum, 
		     int& startLayerNo, int & layerNum, 
		     std::vector<std::string>& labelPaths, 
		     std::vector<std::string>& backgroundPaths, 
		     std::vector<std::string>& adjacentAreaPaths, 
		     std::vector<std::string>& skipAreaPaths, 
		     std::vector<std::string>& adjacentCorrelationPaths, 
		     std::vector<std::string>& skipCorrelationPaths, 
		     std::vector<std::string>& outputPaths, 
		     std::string& xmlPath, std::string& initialFilePath, 
		     std::vector<int>& fwdStartComponentNoSet)
{
  int argNo = 1;
  std::string verboseString = argv[argNo++]; // ARG: verbose level
  sigma_area = atof(argv[argNo++]); // ARG: sigma_area
  sigma_correlation = atof(argv[argNo++]); // ARG: sigma_correlation
  alpha = atof(argv[argNo++]);		  // ARG: alpha
  startSliceNo = atoi(argv[argNo++]);	  // ARG: start slice #
  endSliceNo = atoi(argv[argNo++]);		  // ARG: end slice #
  pileLayerNum = atoi(argv[argNo++]);	  // ARG: pile layer #
  std::string labelTemplate = argv[argNo++];  // ARG: label paths template
  std::string backgroundTemplate = argv[argNo++]; // ARG: bg paths template
  std::string areaTemplate = argv[argNo++];	  // ARG: area paths template
  std::string correlationTemplate = argv[argNo++]; // ARG: correlation paths template
  std::string outputTemplate = argv[argNo++]; // ARG: output paths template
  xmlPath = argv[argNo++];	  // ARG: output xml file path
  initialFilePath = argv[argNo++];
  if (verboseString.compare("INFO") == 0) {
    verbose = VERBOSE_INFO;
  }
  else if (verboseString.compare("HIGH") == 0) {
    verbose = VERBOSE_HIGH;
  }
  else if (verboseString.compare("MEDIUM") == 0) {
    verbose = VERBOSE_MEDIUM;
  }
  else if (verboseString.compare("LOW") == 0) {
    verbose = VERBOSE_LOW;
  }
  else {
    verbose = VERBOSE_NONE;
  }  
  startLayerNo = 0;
  layerNum = endSliceNo - startSliceNo + 1;
  std::string str;
  labelPaths.assign(layerNum, str);
  backgroundPaths.assign(layerNum, str);
  adjacentAreaPaths.assign(layerNum - 1, str);
  skipAreaPaths.assign(layerNum - 2, str);
  adjacentCorrelationPaths.assign(layerNum - 1, str);
  skipCorrelationPaths.assign(layerNum - 2, str);
  outputPaths.assign(layerNum, str);
  for (int sliceNo = startLayerNo; sliceNo < layerNum; sliceNo++) {
    labelPaths[sliceNo] = makePath(labelTemplate, startSliceNo + sliceNo);
    backgroundPaths[sliceNo] = makePath(backgroundTemplate, 
					startSliceNo + sliceNo);
    outputPaths[sliceNo] = makePath(outputTemplate, startSliceNo + sliceNo);
    if (sliceNo < layerNum - 1) {
      adjacentAreaPaths[sliceNo] = 
	makePath(areaTemplate, startSliceNo + sliceNo, 
		 startSliceNo + sliceNo + 1);
      adjacentCorrelationPaths[sliceNo] = 
	makePath(correlationTemplate, startSliceNo + sliceNo, 
		 startSliceNo + sliceNo + 1);
    }
    if (sliceNo < layerNum - 2) {
      skipAreaPaths[sliceNo] = 
	makePath(areaTemplate, startSliceNo + sliceNo, 
		 startSliceNo + sliceNo + 2);
      skipCorrelationPaths[sliceNo] = 
	makePath(correlationTemplate, startSliceNo + sliceNo, 
		 startSliceNo + sliceNo + 2);
    }
  }
  fwdStartComponentNoSet.reserve(argc - argNo);
  while (argNo < argc) {
    int label = atoi(argv[argNo++]);
    fwdStartComponentNoSet.push_back(label);
  }
}



void operation (int argc, char* argv[])
{
  VerboseLevel verbose;
  double sigma_area, sigma_correlation, alpha;
  int startSliceNo, endSliceNo, pileLayerNum, startLayerNo, layerNum;
  std::vector<std::string> labelPaths, backgroundPaths, 
    adjacentAreaPaths, skipAreaPaths, 
    adjacentCorrelationPaths, skipCorrelationPaths, 
    outputPaths;
  std::string xmlPath, initialFilePath;
  std::vector<int> fwdStartComponentNoSet;
  initialization(argc, argv, verbose, sigma_area, sigma_correlation, alpha, 
		 startSliceNo, endSliceNo, pileLayerNum, startLayerNo, layerNum, 
		 labelPaths, backgroundPaths, adjacentAreaPaths, skipAreaPaths, 
		 adjacentCorrelationPaths, skipCorrelationPaths, outputPaths, 
		 xmlPath, initialFilePath, fwdStartComponentNoSet);

  // Generate volume
  clock_t startTime_genVol = clock();
  Volume volume(startSliceNo, startLayerNo, layerNum);
  generateVolume(labelPaths, volume, verbose);
  if (verbose >= VERBOSE_LOW) {
    std::cout << "Generate volume: done." << std::endl;
  }
  clock_t endTime_genVol = clock();

  // Construct graph
  clock_t startTime_constGraph = clock();
  Graph3D graph;
  generateGraph(volume, graph, 
		adjacentCorrelationPaths, skipCorrelationPaths, 
		adjacentAreaPaths, skipAreaPaths, 
		sigma_area, sigma_correlation, alpha, verbose);
  if (verbose >= VERBOSE_LOW) {
    std::cout << "Generate graph: done." << std::endl;
  }
  clock_t endTime_constGraph = clock();

  // Segment
  clock_t startTime_segment = clock();
  SegmentationGraph fwdSegGraph, bwdSegGraph;
  double fwdMaxFlow, bwdMaxFlow;
  // segment(graph, fwdStartComponentNoSet, startLayerNo, layerNum, pileLayerNum, 
	  // fwdSegGraph, bwdSegGraph, fwdMaxFlow, bwdMaxFlow, verbose);
#ifdef GREEDY
  int order = greedy(graph, initialFilePath, fwdStartComponentNoSet, 
  		     startSliceNo, startLayerNo, layerNum, 
  		     fwdSegGraph, bwdSegGraph, 
  		     fwdMaxFlow, bwdMaxFlow, verbose);
  // std::vector<std::vector<int> > startComponentNoSets;
  // std::map<std::vector<int>, SegmentationGraph> fwdSegGraphSets;
  // std::map<std::vector<int>, SegmentationGraph> bwdSegGraphSets;
  // int order = greedy(graph, initialFilePath, 
  // 		     startSliceNo, startLayerNo, layerNum, 
  // 		     verbose);
  // return;
#else
  segment(graph, fwdStartComponentNoSet, startLayerNo, layerNum, 
	  fwdSegGraph, bwdSegGraph, fwdMaxFlow, bwdMaxFlow, verbose);
#endif
  if (verbose >= VERBOSE_LOW) {
    std::cout << "Graphcuts: done." << std::endl;
  }
  clock_t endTime_segment = clock();

  // Group
  clock_t startTime_group = clock();
  SegmentationGraph resSegGraph;
  extractGroup(volume, graph, startLayerNo, fwdSegGraph, bwdSegGraph, 
	       fwdStartComponentNoSet, resSegGraph);
  clock_t endTime_group = clock();

#ifdef DEBUG
  debug_graph_connection(graph, startSliceNo, layerNum);
  return;
#endif

  // write XML
  clock_t startTime_xml = clock();
  writeXML(volume, graph, resSegGraph, startSliceNo, 
	   xmlPath.c_str(), verbose, order);
  if (verbose >= VERBOSE_LOW) {
    std::cout << "XML: done." << std::endl;
  }
  clock_t endTime_xml = clock();

  // Visualize
  clock_t startTime_vis = clock();
  drawRGB(volume, graph, resSegGraph.getCertainTypeSegGraph(SEG_FWD), 
	  resSegGraph.getCertainTypeSegGraph(SEG_BWD), 
	  backgroundPaths, outputPaths, verbose);
  if (verbose >= VERBOSE_LOW) {
    std::cout << "Visualize : done." << std::endl;
  }
  clock_t endTime_vis = clock();

  // Display summary
  if (verbose >= VERBOSE_LOW || verbose == VERBOSE_INFO) {
#ifdef GREEDY
    std::cout << "Order = " << order << std::endl;
#endif
    std::cout << "Forward max flow = " << fwdMaxFlow << std::endl;
    std::cout << "Backward max flow = " << bwdMaxFlow << std::endl;
    std::cout << "Forward cut extracts " << fwdSegGraph.segNodeSet.size() 
	      << " node(s)." << std::endl;
    std::cout << "Backward cut extracts " << bwdSegGraph.segNodeSet.size() 
	      << " node(s)." << std::endl;
    std::cout << "Together extracts " 
	      << (fwdSegGraph.getUnion(bwdSegGraph)).segNodeSet.size() 
	      << " node(s)." << std::endl;
    // std::cout << "Generate volume takes " 
	      // << double(endTime_genVol - startTime_genVol) / CLOCKS_PER_SEC 
	      // << " sec." << std::endl;
    // std::cout << "Construct graph takes " 
	      // << double(endTime_constGraph - startTime_constGraph) / CLOCKS_PER_SEC 
	      // << " sec." << std::endl;
    // std::cout << "Add terminals and segmentation takes "
	      // << double(endTime_segment - startTime_segment) / CLOCKS_PER_SEC 
	      // << " sec." << std::endl;
    // std::cout << "Grouping takes " 
	      // << double(endTime_group - startTime_group) / CLOCKS_PER_SEC 
	      // << " sec." << std::endl;
    // std::cout << "Writing XML takes " 
	      // << double(endTime_xml - startTime_xml) / CLOCKS_PER_SEC 
	      // << " sec." << std::endl;
    // std::cout << "Visualization takes " 
	      // << double(endTime_vis - startTime_vis) / CLOCKS_PER_SEC 
	      // << " sec." << std::endl;
  }
}
